#!/usr/bin/env python

###############################################################################
#
# Use masks to count storm tracks within ocean basins.
#
###############################################################################

#
# MODULES
#
import sys
import os
import re
import glob
import string

#from numarray import *
from pylab import *
#import Numeric

#
# ARGUMENTS
#


# dictionary of (basin, count) and (basin, mask)
past_tccounts = dict()
future_tccounts = dict()

runs = [2, 3, 4, 5, 6, 7, 8, 10, 11, 13, 14, 15, 17, 18]
basins = ['north', 'northatlantic', 'northindian', 'northpacific',
          'south', 'southindian', 'southpacific']

ibasin = dict()
i=0
for basin in basins:
    ibasin[basin] = i
    i=i+1

for basin in basins:
    past_tccounts[basin] = dict()
    future_tccounts[basin] = dict()
    
# loop through runs
for nrun in runs:

    pastfilename = 'basin/tccount_%d_past.txt' % (nrun)
    if os.path.exists(pastfilename):
        for line in open(pastfilename).readlines():

            elems = string.split(line)
            basin = elems[0]
            count = int(elems[1])
        
            past_tccounts[basin][nrun] = count

    futurefilename = 'basin/tccount_%d_future.txt' % (nrun)
    if os.path.exists(futurefilename):
        for line in open(futurefilename).readlines():
            
            elems = string.split(line)
            basin = elems[0]
            count = int(elems[1])
        
            future_tccounts[basin][nrun] = count

#print past_tccounts
#print future_tccounts

# compute means and standard errors
past_means = dict()
future_means = dict()

past_stds = dict()
future_stds = dict()

# plot
#figure(1)
#hold(True)

numruns = len(runs)

for basin in basins:
    ptc = []
    ftc = []
    for nrun in runs:
        ptc.append(past_tccounts[basin][nrun])
        ftc.append(future_tccounts[basin][nrun])

    past_means[basin] = mean(ptc)
    future_means[basin] = mean(ftc)

    past_stds[basin] = std(ptc)
    future_stds[basin] = std(ftc)

for basin in basins:
    
    width = .4
    bar(ibasin[basin] - width*.5, past_means[basin], width*.5,0,'r')
    bar(ibasin[basin], future_means[basin], width*.5,0,'b')

    legend(('Past','Future'))
    title('Per Basin Storm Counts')

    errorbar(ibasin[basin] - width*.25, past_means[basin], past_stds[basin]/sqrt(numruns), None, fmt='k')
    errorbar(ibasin[basin] + width*.25, future_means[basin], future_stds[basin]/sqrt(numruns), None, fmt='k')    

    xlabel('Basin')
    ylabel('Storm Count')

    xticks(arange(len(basins)), basins, rotation=35, fontsize='smaller')


show()
